import sys

sys.path.insert(0, "build/lib.linux-i686-2.5")


import pywcs
import pyfits
import numpy as np

import pywcsgrid


import pylab

def plot_grid(ax, g, draw_grid=True):

    if draw_grid:
        grid_lines1 = [ ax.plot(x, y, "k-") for x, y in g.axis1["grid_lines"]]
        grid_lines2 = [ ax.plot(x, y, "k-") for x, y in g.axis2["grid_lines"]]
        grid_lines = grid_lines1 + grid_lines2
    else:
        grid_lines = []


    ax.set_xticks(g.axis2["ticks_bottom"])
    ax.set_xticklabels(g.axis2["ticklabels_bottom"])

    ax.set_yticks(g.axis1["ticks_left"])
    ax.set_yticklabels(g.axis1["ticklabels_left"])


    ax.xaxis.tick_bottom()
    ax.yaxis.tick_left()
    
    return grid_lines



def plot_fits2(fname):
    f = pyfits.open(fname)
    d, h = f[0].data, f[0].header

    d = d.astype("d")
    d[np.isnan(d)] = -1.

    w = pywcs.WCS(h)

    fig1 = pylab.figure(2, figsize=(12,4))
    fig1.subplots_adjust(wspace=0.3)

    ax1 = pylab.subplot(1,3,1)
    ax2 = pylab.subplot(1,3,2)
    ax3 = pylab.subplot(1,3,3)

    ax1.imshow(d, origin="lower", interpolation="nearest")
    ax2.imshow(d, origin="lower", interpolation="nearest")
    ax3.imshow(d, origin="lower", interpolation="nearest")

    g1 = pywcsgrid.grid(w, h["NAXIS1"], h["NAXIS2"],
                        #label_control=(2,1,2,1)
                        )

    plot_grid(ax1, g1)
    ax1.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax1.set_ylim(-0.5, h["NAXIS2"]-0.5)

    ax1.set_xlabel("Dec.")
    ax1.set_ylabel("R.A.")

    g2 = pywcsgrid.grid(w, h["NAXIS1"], h["NAXIS2"],
                        #label_control=(2,1,2,1),
                        label_density=(12, 6),
                        coord_format = "AA",
                        grid1 = [270, 285, 300],
                        grid2 = [-70, -65],
                        )


    plot_grid(ax2, g2)
    ax2.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax2.set_ylim(-0.5, h["NAXIS2"]-0.5)
    ax2.set_xlabel("Dec.")

    plot_grid(ax3, g1, draw_grid=False)
    ax3.set_xlabel("Dec.")
    ax3.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax3.set_ylim(-0.5, h["NAXIS2"]-0.5)

    pylab.draw()


def plot_fits(fname):
    f = pyfits.open(fname)
    d, h = f[0].data, f[0].header

    d = d.astype("d")
    d[np.isnan(d)] = -1.

    w = pywcs.WCS(h)

    fig1 = pylab.figure(1, figsize=(6,6))
    fig1.subplots_adjust(left=0.2)

    ax1 = pylab.subplot(1,1,1)

    ax1.imshow(d, origin="lower", interpolation="nearest")

    g1 = pywcsgrid.grid(w, h["NAXIS1"], h["NAXIS2"],
                        #label_control=(2,1,2,1)
                        )

    plot_grid(ax1, g1)
    ax1.set_xlim(-0.5, h["NAXIS1"]-0.5)
    ax1.set_ylim(-0.5, h["NAXIS2"]-0.5)

    ax1.set_xlabel("Dec.")
    ax1.set_ylabel("R.A.")


    pylab.draw()



if __name__ == "__main__":

    import sys
    fname = sys.argv[1]

    plot_fits(fname)
    #plot_fits2(fname)
    
    pylab.show()
